'''
Author: Ryan
Email: 409266255@qq.com
Date: 2025-01-28 10:00:00

LastEditors: Ryan
LastEditTime: 2025-09-30 10:53:56
Description: SQLite数据库工具类

FilePath: /qopilot/qopilot/data/database.py
Copyright (c) 2025 by Ryan, All Rights Reserved. 
'''

from enum import Enum
import sqlite3
import pandas as pd
import os
import logging

from qopilot.config import config
from typing import Optional, List, Dict, Any
from contextlib import contextmanager

# 配置日志
logger = logging.getLogger(__name__)


class TableSchema(Enum):
    # 港股列表表结构定义
    HK_STOCKS_LIST_SCHEMA = {
        'code': 'TEXT NOT NULL',
                'name': 'TEXT',
                'date': 'TEXT',
    }

    # 港股历史行情数据表结构定义
    HK_STOCK_HIST_SCHEMA = {
        'date': 'TEXT NOT NULL',
        'open': 'REAL',
        'close': 'REAL',
        'high': 'REAL',
        'low': 'REAL',
        'volume': 'REAL',
        'amount': 'REAL',
        'pct_change': 'REAL',
        'created_at': 'TEXT DEFAULT CURRENT_TIMESTAMP'
    }


class SQLiteManager:
    """SQLite数据库管理类"""

    def __init__(self, db_path: str = None):
        """
        初始化SQLite管理器

        Args:
            db_path: 数据库文件路径，默认为项目根目录下的database文件夹
        """
        if db_path is None:
            # 默认数据库路径
            database_dir = config.DATABASE_DIR
            os.makedirs(database_dir, exist_ok=True)
            db_path = os.path.join(database_dir, 'qopilot.db')

        self.db_path = db_path
        logger.info(f"SQLite数据库路径: {self.db_path}")

    @contextmanager
    def get_connection(self):
        """获取数据库连接的上下文管理器"""
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row  # 使结果可以通过列名访问
            yield conn
        except Exception as e:
            if conn:
                conn.rollback()
            logger.error(f"数据库操作失败: {str(e)}")
            raise
        finally:
            if conn:
                conn.close()

    def create_table(self, table_name: str, columns: Dict[str, str],
                     primary_key: Optional[str] = None,
                     unique_constraints: Optional[List[str]] = None):
        """
        创建数据表

        Args:
            table_name: 表名
            columns: 列定义字典，格式为 {'列名': '数据类型'}
            primary_key: 主键列名
            unique_constraints: 唯一约束列名列表
        """
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()

                # 构建CREATE TABLE语句
                column_definitions = []
                for col_name, col_type in columns.items():
                    column_definitions.append(f"{col_name} {col_type}")

                # 添加主键
                if primary_key:
                    column_definitions.append(f"PRIMARY KEY ({primary_key})")

                # 添加唯一约束
                if unique_constraints:
                    for constraint in unique_constraints:
                        column_definitions.append(f"UNIQUE ({constraint})")

                create_sql = f"""
                CREATE TABLE IF NOT EXISTS {table_name} (
                    {', '.join(column_definitions)}
                )
                """

                logger.info(f"创建表语句: {create_sql}")

                cursor.execute(create_sql)
                conn.commit()
                logger.info(f"成功创建表: {table_name}")

        except Exception as e:
            logger.error(f"创建表失败: {str(e)}")
            raise

    def insert_dataframe(self, df: pd.DataFrame, table_name: str,
                         if_exists: str = 'replace',
                         index: bool = False):
        """
        将DataFrame插入到数据库表中

        Args:
            df: 要插入的DataFrame
            table_name: 目标表名
            if_exists: 如果表存在时的处理方式 ('replace', 'append', 'fail')
            index: 是否包含DataFrame的索引
        """
        try:
            with self.get_connection() as conn:
                df.to_sql(table_name, conn, if_exists=if_exists, index=index)
                conn.commit()
                logger.info(f"成功插入数据到表 {table_name}，共 {len(df)} 条记录")

        except Exception as e:
            logger.error(f"插入数据失败: {str(e)}")
            raise

    def query_data(self, sql: str, params: Optional[tuple] = None) -> List[Dict[str, Any]]:
        """
        查询数据

        Args:
            sql: SQL查询语句
            params: 查询参数

        Returns:
            查询结果列表
        """
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                if params:
                    cursor.execute(sql, params)
                else:
                    cursor.execute(sql)

                # 将结果转换为字典列表
                columns = [description[0] for description in cursor.description]
                results = []
                for row in cursor.fetchall():
                    results.append(dict(zip(columns, row)))

                logger.info(f"查询语句: {sql},params: {params}")
                logger.info(f"查询完成，返回 {len(results)} 条记录")
                return results

        except Exception as e:
            logger.error(f"查询数据失败: {str(e)}")
            raise

    def table_exists(self, table_name: str) -> bool:
        """检查表是否存在"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute("""
                    SELECT name FROM sqlite_master 
                    WHERE type='table' AND name=?
                """, (table_name,))
                return cursor.fetchone() is not None
        except Exception as e:
            logger.error(f"检查表是否存在失败: {str(e)}")
            return False

    def get_table_info(self, table_name: str) -> List[Dict[str, Any]]:
        """获取表结构信息"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute(f"PRAGMA table_info({table_name})")

                columns = [description[0] for description in cursor.description]
                results = []
                for row in cursor.fetchall():
                    results.append(dict(zip(columns, row)))

                return results
        except Exception as e:
            logger.error(f"获取表结构信息失败: {str(e)}")
            return []


class SQLiteManagerHK(SQLiteManager):
    """港股SQLite管理类"""

    def __init__(self):
        super().__init__(db_path=config.DATABASE_PATH_HK)


class SQLiteManagerUS(SQLiteManager):
    """美股SQLite管理类"""

    def __init__(self):
        super().__init__(db_path=config.DATABASE_PATH_US)


if __name__ == "__main__":
    pass
